from __future__ import absolute_import
from __future__ import division
import sys
import os
import argparse
import yaml
import logging
from src import const
from src.segmentation import SegRegion
from src.google_tile import download_tiles_all_count
from src.geo_func import way2tiles
import numpy as np
import pickle

sys.path.append('./src')
from src.inference import MyModel


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('-c', '--conf', help="conf file", default="../geo_server/conf/server.yml", type=argparse.FileType('r'))
    parser.add_argument('-r', '--region', type=str, default='shanghai')
    parser.add_argument('--china-convert', type=bool, default=True)
    parser.add_argument('--dir', default='data/tiles')
    parser.add_argument('--save-file', default='data/singe_feature.pkl')
    parser.add_argument('--num-classes', default=6, type=int)
    return parser.parse_args()


def main(args):
    conf = yaml.load(args.conf)
    region = getattr(const, args.region.upper() + '_REGION')
    way = region.to_way()
    print(way.left_up_corner)
    print(way.right_down_corner)
    tiles = way2tiles(way, args.china_convert)
    # for tile in tiles:
    #     print(tile[0], tile[1], tile[2])
    logging.info(len(tiles))
    download_tiles_all_count(tiles, args.dir)

    if os.path.exists(args.save_file):
        logging.warning('Load feature_dict from file.')
        with open(args.save_file, 'rb') as f:
            feature_dict = pickle.load(f)
    else:
        feature_dict = {}

    model = MyModel(conf, ext='jpg')
    for i, tile in enumerate(tiles):
        filename = '%d_%d_%d.jpg' % (tile[0], tile[1], tile[2])
        if filename in feature_dict:
            continue
        filepath_name = os.path.join(args.dir, filename)
        if os.path.exists(filepath_name):
            one_feature, _ = model.detect_one_img(filepath_name)
        feature_dict[filename] = one_feature
        if (i + 1) % 100 == 0:
            logging.info('%d / %d. Save file.' % (i + 1, len(tiles)))
            with open(args.save_file, 'wb') as f:
                pickle.dump(feature_dict, f)

    with open(args.save_file, 'wb') as f:
                pickle.dump(feature_dict, f)


if __name__ == '__main__':
    args = parse_args()
    logging.basicConfig(level=logging.INFO)
    main(args)
